package cn.hinglo.common.interceptors;

import cn.hinglo.common.entity.LimitationInfo;
import cn.hinglo.common.enums.LimitCacheTypeEnum;
import cn.hinglo.common.enums.LimitRenderTypeEnum;
import cn.hinglo.common.util.ResultKit;
import cn.hutool.core.util.NumberUtil;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.util.concurrent.RateLimiter;
import com.jfinal.core.Controller;
import com.jfinal.plugin.redis.Cache;
import com.jfinal.plugin.redis.Redis;
import lombok.Data;

import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;


/**
 * @author HingLo
 * @createTime: 2018/9/25 10:03
 * @description: 限流管理工具
 */
@Data
public final class RateLimitManager {

    /***
     * 是否启用限流
     */
    public static Boolean enable = true;
    /***
     * 缓存名称（主要是对于在Redis与ehcache中缓存到制定名称下）
     */
    private static String CACHE_NAME = "HingLoRateLimit";

    /***
     * 缓存类型
     */
    private static LimitCacheTypeEnum cacheType = LimitCacheTypeEnum.MEMORY;

    /***
     * 默认返回的Json数据类型
     */
    private static LimitRenderTypeEnum renderType = LimitRenderTypeEnum.JSON;


    /**
     * 自身单例对象
     */
    private static final RateLimitManager me = new RateLimitManager();


    /**
     * 创建一个线程安全的map类型的令牌桶,使用内存缓存限流信息
     */
    private LoadingCache<String, RateLimiter> requestRateLimiterMap = CacheBuilder.newBuilder().expireAfterAccess(60, TimeUnit.SECONDS)
            .build(new CacheLoader<String, RateLimiter>() {

                @Override
                public RateLimiter load(String key) {
                    return null;
                }
            });

    /**
     * 构造函数,默认要初始化配置信息
     */
    private RateLimitManager() {
    }

    /**
     * 单例模式
     *
     * @return 返回单例对象
     */
    public static RateLimitManager me() {
        return me;
    }


    /**
     * 并发限流
     *
     * @param limitationInfo 请求的目标
     * @param limitationInfo 限流信息
     * @param controller     响应信息
     * @return 返回限流结果
     */
    public boolean rateLimit(LimitationInfo limitationInfo, Controller controller) {
        // 限流信息处理
        boolean res = limit(limitationInfo.getTarget(), limitationInfo.getRate(), limitationInfo.getSecond());
        if (!res) {
            this.render(limitationInfo, controller);
        }
        return res;
    }


    /**
     * 初始化一个限流的目标
     *
     * @param target 目标对象,ip,url,user,request等
     * @param rate   每秒通过的路由
     * @return 返回限流对象
     */
    private RateLimiter initRateLimiter(String target, double rate) {
        // 保留2位小数就好
        rate = NumberUtil.round(rate, 2).doubleValue();

        RateLimiter limiter = getLimiter(target);
        if (limiter == null) {
            limiter = RateLimiter.create(rate);
            setLimiter(target, limiter);
            return limiter;
        }

        if (limiter.getRate() == rate) {
            return limiter;
        }
        limiter.setRate(rate);
        setLimiter(target, limiter);
        return limiter;
    }

    /**
     * 通过指定的可以来返回限流对象
     *
     * @param target 目标对象
     * @return 返回限流对象, 默认使用内存作为缓存
     */
    private RateLimiter getLimiter(String target) {
        try {
            return requestRateLimiterMap.get(target);
        } catch (ExecutionException | RuntimeException e) {
            return null;
        }
    }

    /**
     * 添加或者更新限流信息
     *
     * @param target  目标
     * @param limiter 限流对象
     */
    private void setLimiter(String target, RateLimiter limiter) {
        requestRateLimiterMap.put(target, limiter);
    }

    /**
     * 限流方法
     *
     * @param target 目标
     * @param rate   限流次数
     * @param second 限流时间
     * @return 返回结果
     */
    private boolean limit(String target, int rate, int second) {
        // 使用内存方案
        if (cacheType == LimitCacheTypeEnum.MEMORY) {
            RateLimiter rateLimiter = initRateLimiter(target, rate * 1.0 / second);
            return rateLimiter.tryAcquire();
        } else if (cacheType == LimitCacheTypeEnum.REDIS) {
            // 下面是使用redis方式限流
            String val = redisGetLimit(target);
            if (val == null) {
                redisSetLimit(target, 1, second);
                return true;
            } else {
                int count = Integer.parseInt(val);
                if (count < rate) {
                    // 请求次数+1
                    redisSetLimit(target, count + 1, -1);
                    return true;
                }
                return false;
            }
        } else {

            // todo EhCache 待实现
            return true;
        }
    }

    /**
     * redis限流方式的存入
     *
     * @param target 目标
     * @param rate   次数
     * @param second 指定的时间
     */
    private void redisSetLimit(String target, int rate, int second) {
        Cache cache = Redis.use(CACHE_NAME);
        // 如果second=-1 表示是更新
        if (second == -1) {
            cache.set(target, String.valueOf(rate));
            // 剩余时间
            long millis = cache.ttl(target);
            // 更新有效期的时间
            cache.pexpire(target, millis);
        } else {
            cache.setex(target, second, String.valueOf(rate));
        }
    }

    /**
     * 通过目标 返回限流信息
     *
     * @param target 目标对象
     * @return 返回结果
     */
    private String redisGetLimit(String target) {
        Cache cache = Redis.use(CACHE_NAME);
        return cache.get(target);
    }


    /***
     * 根据类型返回结果
     * @param limitationInfo 限流信息
     * @param controller 请求信息
     */
    private void render(LimitationInfo limitationInfo, Controller controller) {
        if (renderType == LimitRenderTypeEnum.JSON) {
            // 如果是返回json格式的数据
            controller.renderJson(ResultKit.error(501, limitationInfo.getContent()));
        } else if (renderType == LimitRenderTypeEnum.PAGE) {
            // 返回页面
            controller.render(limitationInfo.getContent());
        } else if (renderType == LimitRenderTypeEnum.REDIRECT) {
            // 重定向处理
            controller.redirect(limitationInfo.getContent());
        }
    }

}
